#!/bin/bash

# K8s离线安装脚本 - 增强版
# 支持三台服务器的集群部署
# 包含详细的错误检查和日志输出

set -e

# 配置变量
LOG_FILE="/var/log/k8s-install.log"
INSTALL_DIR="/root/k8s-offline-packages"
NODE_TYPE=""  # master 或 worker
MASTER_IP=""
CLUSTER_TOKEN=""
CLUSTER_CERT_HASH=""

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log() {
    local level=$1
    shift
    local message="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    case $level in
        "INFO")
            echo -e "${BLUE}[INFO]${NC} $message" | tee -a $LOG_FILE
            ;;
        "SUCCESS")
            echo -e "${GREEN}[SUCCESS]${NC} $message" | tee -a $LOG_FILE
            ;;
        "WARNING")
            echo -e "${YELLOW}[WARNING]${NC} $message" | tee -a $LOG_FILE
            ;;
        "ERROR")
            echo -e "${RED}[ERROR]${NC} $message" | tee -a $LOG_FILE
            ;;
    esac
}

# 错误处理函数
error_exit() {
    log "ERROR" "$1"
    log "ERROR" "安装失败，请检查日志: $LOG_FILE"
    exit 1
}

# 检查函数
check_root() {
    if [ "$EUID" -ne 0 ]; then
        error_exit "请使用root用户运行此脚本"
    fi
}

check_os() {
    log "INFO" "检查操作系统..."
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        if [[ "$ID" == "ubuntu" ]]; then
            log "SUCCESS" "操作系统检查通过: $PRETTY_NAME"
        else
            log "WARNING" "当前系统不是Ubuntu，可能不兼容"
        fi
    else
        log "WARNING" "无法确定操作系统版本"
    fi
}

check_memory() {
    log "INFO" "检查内存..."
    local mem_gb=$(free -g | awk '/^Mem:/{print $2}')
    if [ $mem_gb -lt 2 ]; then
        error_exit "内存不足，至少需要2GB内存，当前: ${mem_gb}GB"
    fi
    log "SUCCESS" "内存检查通过: ${mem_gb}GB"
}

check_disk() {
    log "INFO" "检查磁盘空间..."
    local disk_gb=$(df / | awk 'NR==2{print int($4/1024/1024)}')
    if [ $disk_gb -lt 20 ]; then
        error_exit "磁盘空间不足，至少需要20GB，当前可用: ${disk_gb}GB"
    fi
    log "SUCCESS" "磁盘空间检查通过: ${disk_gb}GB可用"
}

# 系统准备
prepare_system() {
    log "INFO" "准备系统环境..."
    
    # 更新包管理器
    log "INFO" "更新包管理器..."
    apt-get update
    
    # 安装必要工具
    log "INFO" "安装必要工具..."
    apt-get install -y curl wget gnupg lsb-release ca-certificates
    
    # 关闭防火墙
    log "INFO" "关闭防火墙..."
    ufw disable 2>/dev/null || true
    
    # 关闭swap
    log "INFO" "关闭swap..."
    swapoff -a
    sed -i '/ swap / s/^\(.*\)$/#\1/g' /etc/fstab
    
    # 配置内核参数
    log "INFO" "配置内核参数..."
    cat > /etc/sysctl.d/k8s.conf << EOF
net.bridge.bridge-nf-call-ip6tables = 1
net.bridge.bridge-nf-call-iptables = 1
net.ipv4.ip_forward = 1
vm.swappiness = 0
EOF
    sysctl --system
    
    # 加载br_netfilter模块
    modprobe br_netfilter
    echo 'br_netfilter' > /etc/modules-load.d/k8s.conf
    
    # 安装Docker（如果未安装）
    if ! command -v docker &> /dev/null; then
        log "INFO" "安装Docker..."
        curl -fsSL https://get.docker.com -o get-docker.sh
        sh get-docker.sh
        systemctl enable docker
        systemctl start docker
    fi
    
    log "SUCCESS" "系统环境准备完成"
}

# 安装containerd
install_containerd() {
    log "INFO" "安装containerd..."
    
    if [ ! -f "$INSTALL_DIR/packages/containerd-"*.tar.gz ]; then
        error_exit "未找到containerd安装包"
    fi
    
    # 解压containerd
    tar -xzf $INSTALL_DIR/packages/containerd-*.tar.gz -C /usr/local/
    
    # 创建systemd服务文件
    cat > /etc/systemd/system/containerd.service << EOF
[Unit]
Description=containerd container runtime
Documentation=https://containerd.io
After=network.target local-fs.target

[Service]
ExecStartPre=-/sbin/modprobe overlay
ExecStart=/usr/local/bin/containerd
Type=notify
Delegate=yes
KillMode=process
Restart=always
RestartSec=5
LimitNPROC=infinity
LimitCORE=infinity
LimitNOFILE=infinity
TasksMax=infinity
OOMScoreAdjust=-999

[Install]
WantedBy=multi-user.target
EOF
    
    # 配置containerd
    mkdir -p /etc/containerd
    containerd config default > /etc/containerd/config.toml
    
    # 启动containerd
    systemctl daemon-reload
    systemctl enable containerd
    systemctl start containerd
    
    # 验证containerd
    if systemctl is-active --quiet containerd; then
        log "SUCCESS" "containerd安装成功"
    else
        error_exit "containerd启动失败"
    fi
}

# 安装runc
install_runc() {
    log "INFO" "安装runc..."
    
    if [ ! -f "$INSTALL_DIR/binaries/runc" ]; then
        error_exit "未找到runc二进制文件"
    fi
    
    cp $INSTALL_DIR/binaries/runc /usr/local/bin/
    chmod +x /usr/local/bin/runc
    
    if command -v runc &> /dev/null; then
        log "SUCCESS" "runc安装成功"
    else
        error_exit "runc安装失败"
    fi
}

# 安装CNI插件
install_cni() {
    log "INFO" "安装CNI插件..."
    
    if [ ! -f "$INSTALL_DIR/packages/cni-plugins-"*.tgz ]; then
        error_exit "未找到CNI插件安装包"
    fi
    
    mkdir -p /opt/cni/bin
    tar -xzf $INSTALL_DIR/packages/cni-plugins-*.tgz -C /opt/cni/bin/
    
    if [ -d "/opt/cni/bin" ] && [ "$(ls /opt/cni/bin | wc -l)" -gt 0 ]; then
        log "SUCCESS" "CNI插件安装成功"
    else
        error_exit "CNI插件安装失败"
    fi
}

# 安装crictl
install_crictl() {
    log "INFO" "安装crictl..."
    
    if [ ! -f "$INSTALL_DIR/packages/crictl-"*.tar.gz ]; then
        error_exit "未找到crictl安装包"
    fi
    
    tar -xzf $INSTALL_DIR/packages/crictl-*.tar.gz -C /usr/local/bin/
    chmod +x /usr/local/bin/crictl
    
    if command -v crictl &> /dev/null; then
        log "SUCCESS" "crictl安装成功"
    else
        error_exit "crictl安装失败"
    fi
}

# 安装K8s二进制文件
install_k8s_binaries() {
    log "INFO" "安装K8s二进制文件..."
    
    # 首先尝试安装DEB包
    if [ -d "$INSTALL_DIR/packages/deb" ] && [ "$(ls $INSTALL_DIR/packages/deb/*.deb 2>/dev/null | wc -l)" -gt 0 ]; then
        log "INFO" "安装K8s DEB包..."
        dpkg -i $INSTALL_DIR/packages/deb/*.deb
        apt-get install -f -y  # 修复依赖关系
    else
        # 如果没有DEB包，使用二进制文件
        log "INFO" "使用二进制文件安装K8s..."
        local binaries=("kubeadm" "kubelet" "kubectl")
        
        for binary in "${binaries[@]}"; do
            if [ ! -f "$INSTALL_DIR/binaries/$binary" ]; then
                error_exit "未找到$binary二进制文件"
            fi
            
            cp $INSTALL_DIR/binaries/$binary /usr/local/bin/
            chmod +x /usr/local/bin/$binary
            
            if command -v $binary &> /dev/null; then
                log "SUCCESS" "$binary安装成功"
            else
                error_exit "$binary安装失败"
            fi
        done
    fi
}

# 安装Helm
install_helm() {
    log "INFO" "安装Helm..."
    
    if [ ! -f "$INSTALL_DIR/packages/helm-"*.tar.gz ]; then
        log "WARNING" "未找到Helm安装包，跳过安装"
        return
    fi
    
    tar -xzf $INSTALL_DIR/packages/helm-*.tar.gz
    mv linux-amd64/helm /usr/local/bin/
    chmod +x /usr/local/bin/helm
    rm -rf linux-amd64
    
    if command -v helm &> /dev/null; then
        log "SUCCESS" "Helm安装成功"
    else
        log "WARNING" "Helm安装失败，但不影响K8s核心功能"
    fi
}

# 加载Docker镜像
load_images() {
    log "INFO" "加载Docker镜像..."
    
    if [ ! -d "$INSTALL_DIR/images" ]; then
        error_exit "未找到镜像目录"
    fi
    
    local image_count=0
    for image in $INSTALL_DIR/images/*.tar; do
        if [ -f "$image" ]; then
            log "INFO" "加载镜像: $(basename $image)"
            docker load < "$image"
            ((image_count++))
        fi
    done
    
    if [ $image_count -gt 0 ]; then
        log "SUCCESS" "成功加载 $image_count 个镜像"
    else
        error_exit "未找到任何镜像文件"
    fi
}

# 配置kubelet
configure_kubelet() {
    log "INFO" "配置kubelet..."
    
    # 创建kubelet配置目录
    mkdir -p /etc/systemd/system/kubelet.service.d
    
    # 创建kubelet服务文件
    cat > /etc/systemd/system/kubelet.service.d/10-kubeadm.conf << EOF
[Service]
Environment="KUBELET_KUBECONFIG_ARGS=--bootstrap-kubeconfig=/etc/kubernetes/bootstrap-kubelet.conf --kubeconfig=/etc/kubernetes/kubelet.conf"
Environment="KUBELET_CONFIG_ARGS=--config=/var/lib/kubelet/config.yaml"
EnvironmentFile=-/var/lib/kubelet/kubeadm-flags.env
ExecStart=
ExecStart=/usr/local/bin/kubelet \$KUBELET_KUBECONFIG_ARGS \$KUBELET_CONFIG_ARGS \$KUBELET_KUBEADM_ARGS \$KUBELET_EXTRA_ARGS
EOF
    
    systemctl daemon-reload
    systemctl enable kubelet
    
    log "SUCCESS" "kubelet配置完成"
}

# 初始化集群（仅master节点）
init_cluster() {
    if [ "$NODE_TYPE" != "master" ]; then
        return
    fi
    
    log "INFO" "初始化K8s集群..."
    
    # 初始化集群
    kubeadm init --pod-network-cidr=10.244.0.0/16 --apiserver-advertise-address=$MASTER_IP
    
    # 配置kubectl
    mkdir -p $HOME/.kube
    cp -i /etc/kubernetes/admin.conf $HOME/.kube/config
    chown $(id -u):$(id -g) $HOME/.kube/config
    
    # 保存join命令
    kubeadm token create --print-join-command > /root/kubeadm-join-command.txt
    
    log "SUCCESS" "集群初始化完成"
    log "INFO" "Join命令已保存到: /root/kubeadm-join-command.txt"
}

# 加入集群（仅worker节点）
join_cluster() {
    if [ "$NODE_TYPE" != "worker" ]; then
        return
    fi
    
    if [ -z "$CLUSTER_TOKEN" ] || [ -z "$CLUSTER_CERT_HASH" ]; then
        error_exit "Worker节点需要提供集群token和证书hash"
    fi
    
    log "INFO" "加入K8s集群..."
    
    kubeadm join $MASTER_IP:6443 --token $CLUSTER_TOKEN --discovery-token-ca-cert-hash sha256:$CLUSTER_CERT_HASH
    
    log "SUCCESS" "节点已加入集群"
}

# 主函数
main() {
    log "INFO" "开始K8s离线安装..."
    log "INFO" "安装日志: $LOG_FILE"
    
    # 检查环境
    check_root
    check_os
    check_memory
    check_disk
    
    # 准备系统
    prepare_system
    
    # 安装组件
    install_containerd
    install_runc
    install_cni
    install_crictl
    install_k8s_binaries
    install_helm
    
    # 加载镜像
    load_images
    
    # 配置kubelet
    configure_kubelet
    
    # 根据节点类型执行相应操作
    if [ "$NODE_TYPE" = "master" ]; then
        init_cluster
    elif [ "$NODE_TYPE" = "worker" ]; then
        join_cluster
    fi
    
    log "SUCCESS" "K8s安装完成！"
    log "INFO" "请运行检查脚本验证安装: ./check_k8s.sh"
}

# 显示使用说明
show_usage() {
    echo "K8s离线安装脚本"
    echo ""
    echo "用法:"
    echo "  $0 --node-type=master --master-ip=<IP>     # 安装master节点"
    echo "  $0 --node-type=worker --master-ip=<IP> --token=<TOKEN> --cert-hash=<HASH>  # 安装worker节点"
    echo ""
    echo "参数:"
    echo "  --node-type     节点类型 (master/worker)"
    echo "  --master-ip     Master节点IP地址"
    echo "  --token         加入集群的token (仅worker节点需要)"
    echo "  --cert-hash     集群证书hash (仅worker节点需要)"
    echo ""
    echo "示例:"
    echo "  $0 --node-type=master --master-ip=192.168.1.10"
    echo "  $0 --node-type=worker --master-ip=192.168.1.10 --token=abc123 --cert-hash=def456"
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        --node-type=*)
            NODE_TYPE="${1#*=}"
            shift
            ;;
        --master-ip=*)
            MASTER_IP="${1#*=}"
            shift
            ;;
        --token=*)
            CLUSTER_TOKEN="${1#*=}"
            shift
            ;;
        --cert-hash=*)
            CLUSTER_CERT_HASH="${1#*=}"
            shift
            ;;
        --help)
            show_usage
            exit 0
            ;;
        *)
            echo "未知参数: $1"
            show_usage
            exit 1
            ;;
    esac
done

# 验证参数
if [ -z "$NODE_TYPE" ]; then
    echo "错误: 必须指定节点类型"
    show_usage
    exit 1
fi

if [ "$NODE_TYPE" != "master" ] && [ "$NODE_TYPE" != "worker" ]; then
    echo "错误: 节点类型必须是master或worker"
    exit 1
fi

if [ -z "$MASTER_IP" ]; then
    echo "错误: 必须指定master IP"
    exit 1
fi

# 运行主函数
main
